package solutions.jianzhi

import solutions.lib.TreeNode

/**
 * @author lizhidong
 * <a href="https://leetcode.cn/problems/er-cha-sou-suo-shu-de-di-kda-jie-dian-lcof/">剑指 Offer 54. 二叉搜索树的第k大节点</a>
 */
private class Solution54 {
    fun kthLargest(root: TreeNode?, k: Int): Int {
        return findNext(root, k)?.second?.`val` ?: -1
    }

    private fun findNext(root: TreeNode?, target: Int): Pair<Int, TreeNode?>? {
        if (root == null) {
            return null
        }
        val rp = findNext(root.right, target)
        if (rp?.second != null) {
            return rp
        }
        val remain = target - (rp?.first ?: 0)
        if (remain == 1) {
            return Pair(0, root)
        }
        val lp = findNext(root.left, remain - 1)
        return if (lp?.second != null) {
            lp
        } else {
            Pair((rp?.first ?: 0) + 1 + (lp?.first ?: 0), null)
        }
    }
}

fun main() {
    println(
        Solution54().kthLargest(
            TreeNode(
                5,
                TreeNode(
                    3,
                    TreeNode(
                        2,
                        TreeNode(1),
                        null
                    ),
                    TreeNode(4)
                ),
                TreeNode(
                    6
                )
            ), 3
        )
    )
}